import { cpus } from 'os'
import { name, version } from '../../package.json'
import { FastWorkerItem, FastWorkerItemStatusError } from '../manager/worker'
import { callbackMap, stringifyCallbackParameterTransferSymbolKey } from '../transfer'
import { _import, _is_esm } from '../utils'

const awaitTimeout = (timeout?: number) => {
  return new Promise<void>(resolve => {
    setTimeout(resolve, timeout)
  })
}

type HookFnInvokeMethod = (
  moduleId: string,
  method: string,
  args: any[],
  timeout?: number | null,
  retry?: number,
  retryInterval?: number
) => Promise<any>
export type HookFn = (invokeMethod: HookFnInvokeMethod, isMainThread: boolean) => void

const doOnHookMain = async (moduleId: string, method: string, args: any[]) => {
  const module = await _import(moduleId)
  const methodObject = module?.[method]
  if (typeof methodObject !== 'function') {
    throw new Error(`Method '${method}' in module '${moduleId}' is not a function.`)
  }
  return methodObject(...args)
}

const fastWorkSymbol = Symbol.for(`${name}@${version}`)
export interface WorkerPool {
  invokeMethod: (
    moduleId: string,
    method: string,
    args: any[],
    timeout?: number | null,
    retry?: number,
    retryInterval?: number,
    onlyExecInWorker?: boolean | undefined
  ) => Promise<any>
  onInit: (fn: HookFn, onlyExecInWorker?: boolean) => void
  offInit: (fn: HookFn) => void
  every: (fn: HookFn, onlyExecInWorker?: boolean) => void
  ref: () => void
  unref: () => void
  [fastWorkSymbol]: any
}

export const initWorkersPool = (): WorkerPool => {
  let worker: WorkerPool = (globalThis as any)[fastWorkSymbol]
  if (worker && worker[fastWorkSymbol] === fastWorkSymbol) {
    return worker
  }

  const maxThreads = Math.max(2, cpus().length - 1)
  const defaultThreads = Math.max(1, Math.trunc(maxThreads / 1.75))
  const workers: FastWorkerItem[] = []
  const initFns: HookFn[] = []

  let mainRunnings = 0
  let _k = 0

  const createWorker = (index = workers.length) => {
    let isRemoved = false

    const worker = new FastWorkerItem({
      FAST_WORKER_LOWER_LEVEL: index >= defaultThreads ? 'true' : ''
    })
    const remove = () => {
      if (isRemoved) return
      isRemoved = true

      const idx = workers.indexOf(worker)
      if (idx >= 0) {
        workers.splice(idx, 1)
      }
    }

    worker.whenClosed().finally(remove)
    worker.whenReady().catch(remove)

    if (initFns.length > 0) {
      const invoke = doOnHookWorker.bind(worker)
      for (const fn of initFns) {
        fn(invoke, false)
      }
    }

    return worker
  }

  const addWorker = () => {
    if (workers.length >= maxThreads) return

    const worker = createWorker()
    workers.push(worker)
    return worker
  }

  const getPriorityWorker = () => {
    let invokedWorker: FastWorkerItem | undefined
    let currentRunnings = Infinity

    for (let i = 0; i < workers.length; ++i) {
      let worker = workers[i]
      if (worker.closed) {
        worker = workers[i] = createWorker(i)
      }
      if (!worker.ready) continue

      const { runnings } = worker
      if (runnings === 0) {
        currentRunnings = 0
        invokedWorker = worker
        break
      } else if (runnings < currentRunnings) {
        currentRunnings = runnings
        invokedWorker = worker
      }
    }

    if (invokedWorker && invokedWorker.runnings >= mainRunnings + 2) {
      invokedWorker = addWorker() ?? invokedWorker
    }

    if (!invokedWorker) {
      const availWorkers = workers.filter(worker => !worker.closed)
      invokedWorker = availWorkers[Math.trunc(Math.random() * availWorkers.length)] ?? workers[0]
    }

    return invokedWorker
  }

  const resolveArgs = (args: any[]) => {
    const noCallbackArgs: any[] = []
    const callbackArgs: (string | number)[] = []

    for (let i = 0, length = args.length; i < length; ++i) {
      const arg = args[i]
      if (typeof arg === 'function') {
        if (arg[stringifyCallbackParameterTransferSymbolKey]) {
          callbackArgs[i] = arg.toString()
        } else {
          const callbackKey = (callbackArgs[i] = ++_k)
          callbackMap.set(callbackKey, arg)
        }
      } else {
        noCallbackArgs[i] = arg
      }
    }

    return [noCallbackArgs, callbackArgs] as const
  }

  const cleanupCallbackArgs = (callbackArgs: (string | number)[]) => {
    for (const v of callbackArgs) {
      if (typeof v === 'number') {
        const asyncCallbackState = callbackMap.get(v)
        if (asyncCallbackState) {
          callbackMap.delete(v)
        }
      }
    }
  }

  const doOnHookWorker = async function (
    this: FastWorkerItem,
    moduleId: string,
    method: string,
    args: any[],
    timeout: number | null = 5000,
    retry: number = 0,
    retryInterval?: number
  ) {
    const [noCallbackArgs, callbackArgs] = resolveArgs(args)
    let count = Math.max(0, Math.trunc(retry)) + 1
    let err: any

    for (let i = 0; i < count; ) {
      try {
        let result = await this.invokeMethod(moduleId, method, noCallbackArgs, callbackArgs, timeout)
        return result
      } catch (error) {
        err = error

        if (++i < count) {
          await awaitTimeout(retryInterval)
        }
      }
    }

    if (callbackArgs!) {
      cleanupCallbackArgs(callbackArgs)
    }

    throw err
  }

  const _invokeMethod = (
    moduleId: string,
    method: string,
    rawArgs: any[],
    timeout: number | null,
    onlyExecInWorker: boolean | undefined,
    getResolvedArgs: () => readonly [any[], (string | number)[]]
  ) => {
    let worker = getPriorityWorker()

    if (!onlyExecInWorker) {
      if (worker.closed || !worker.ready || worker.runnings >= mainRunnings + 2) {
        ++mainRunnings
        return _import(moduleId)
          .then(module => {
            const methodObject = module?.[method]
            if (typeof methodObject !== 'function') {
              throw new Error(`Method '${method}' in module '${moduleId}' is not a function.`)
            }

            return methodObject(...rawArgs)
          })
          .finally(() => {
            --mainRunnings
          })
      }
    }

    const [noCallbackArgs, callbackArgs] = getResolvedArgs()
    return worker.invokeMethod(moduleId, method, noCallbackArgs, callbackArgs, timeout)
  }

  const invokeMethod = async (
    moduleId: string,
    method: string,
    args: any[],
    timeout: number | null = 5000,
    retry = 0,
    retryInterval?: number,
    onlyExecInWorker: boolean | undefined = _is_esm ? undefined : moduleId.endsWith('.mjs')
  ) => {
    let count = Math.max(0, Math.trunc(retry)) + 1
    let err: any

    let noCallbackArgs: any[]
    let callbackArgs: (string | number)[]

    const getResolvedArgs = () => {
      if (callbackArgs) return [noCallbackArgs, callbackArgs] as const
      const result = resolveArgs(args)
      ;[noCallbackArgs, callbackArgs] = result
      return result
    }

    for (let i = 0; i < count; ) {
      try {
        let result = await _invokeMethod(moduleId, method, args, timeout, onlyExecInWorker, getResolvedArgs)
        return result
      } catch (error) {
        err = error
        if (error instanceof FastWorkerItemStatusError) {
          ++count
        }

        if (++i < count) {
          await awaitTimeout(retryInterval)
        }
      }
    }

    if (callbackArgs!) {
      cleanupCallbackArgs(callbackArgs)
    }

    throw err
  }

  let ref = 0
  let timer: any

  return ((globalThis as any)[fastWorkSymbol] = {
    invokeMethod,
    onInit: (fn: HookFn, onlyExecInWorker = false) => {
      for (const worker of workers) {
        fn(doOnHookWorker.bind(worker), false)
      }
      initFns.push(fn)
      if (!onlyExecInWorker) {
        fn(doOnHookMain, true)
      }
    },
    offInit: (fn: HookFn) => {
      const idx = initFns.indexOf(fn)
      if (idx >= 0) {
        initFns.splice(idx, 1)
      }
    },
    every: (fn: HookFn, onlyExecInWorker = false) => {
      for (const worker of workers) {
        if (!worker.closed) {
          fn(doOnHookWorker.bind(worker), false)
        }
      }
      if (!onlyExecInWorker) {
        fn(doOnHookMain, true)
      }
    },
    ref: () => {
      ++ref
      clearTimeout(timer)
      if (workers.length < defaultThreads) {
        for (let i = 0; i < defaultThreads; ++i) {
          addWorker()
        }
      }
    },
    unref: () => {
      --ref
      clearTimeout(timer)
      if (ref === 0) {
        timer = setTimeout(
          () => {
            timer = undefined
            if (ref === 0) {
              for (const worker of workers) {
                worker.terminate()
              }
              workers.splice(0, workers.length)
            }
          },
          timer ? 1000 : 100
        )
      }
    },
    [fastWorkSymbol]: fastWorkSymbol
  })
}
